#!/usr/bin/env python3
"""
Point Cloud Painting Utility

This module provides paintPcd function for recoloring segpointcloud in unifiedData
according to allowed_seg_colors.json's indices2obj for the current embodiment.
"""

import numpy as np
from termcolor import colored
from semPcd import semPcd

def paintPcd(unifiedData: dict, embodiment: str, surpress_print: bool = False, seg_source: str = "traverse") -> dict:
    """
    Recolor segpointcloud in unifiedData by mapping its current RGB values (from embodiment-specific colors)
    to target RGB values (from top-level indices2obj), and store the result in a new field 'paintpointcloud'.
    """
    assert 'segpointcloud' in unifiedData and unifiedData['segpointcloud'] is not None, "paintPcd: No segpointcloud to paint."
    segpcd = unifiedData['segpointcloud']
    assert segpcd.shape[0] > 0, "paintPcd: segpointcloud is empty."
    semantic_indices, top_indices2obj = semPcd(unifiedData, embodiment, surpress_print=False, return_colors=True, seg_source=seg_source)
    assert semantic_indices, "paintPcd: No semantic parts found."
    new_segpcd = segpcd.copy()
    painted = 0
    for semantic_part, point_indices in semantic_indices.items():
        top_val = top_indices2obj.get(semantic_part)
        assert top_val is not None, f"paintPcd: No target color found for {semantic_part}"
        target_color = top_val[0:3] if isinstance(top_val, list) else None
        assert target_color is not None, f"paintPcd: Invalid target color for {semantic_part}"
        if point_indices:
            new_segpcd[point_indices, 3:6] = target_color
            painted += len(point_indices)
            if not surpress_print:
                print(f"    🎨 PAINT | {semantic_part}: {len(point_indices)} points → {target_color}")
    unifiedData['paintpointcloud'] = new_segpcd
    if not surpress_print:
        print(f"    🎨 PAINT | Total {painted} points recolored (output: paintpointcloud)")
    return unifiedData 